import torch
# 假设我们有两个相同的骰子，但是其中的一个是公平的，每个点数出现的概率相同；另一个骰子则被做了手脚，数字 6 出现的概率为 80%，
# 而数字 1-5 出现的概率都为 4%。如果我给你一个 15 次投掷骰子的序列，你能预测出我每次投掷用的是哪一枚骰子吗
def neg_log_likelihood(self, rolls, states):
    """Compute neg log-likelihood for a given sequence.

    Input: 
    rolls: numpy array, dim [1, n_rolls]. Integer 0-5 showing value on dice.
    states: numpy array, dim [1, n_rolls]. Integer 0, 1. 0 if dice is fair.
    """
    """个方法做了三件主要的事情：
    1）将骰子的值映射到一个似然函数上 
    2）计算对数似然项的分子 
    3）计算对数似然项的分母
    """

    loglikelihoods = self._data_to_likelihood(rolls)
    states = torch.LongTensor(states)
    sequence_loglik = self._compute_likelihood_numerator(loglikelihoods, states)
    denominator = self._compute_likelihood_denominator(loglikelihoods)
    return denominator - sequence_loglik

def _data_to_likelihood(self, rolls):
    """
    创建一个维度为 6*2 的矩阵，其中第一列是用公平骰子掷到 1-6 的似然度，
    第二列是用有偏骰子掷到 1-6 的似然度
    Converts a numpy array of rolls (integers) to log-likelihood.
    self.loglikeligood is a matrix of 6 x 2 in our case. 
    Input is one [1, n_rolls]
    """
    log_likelihoods = self.loglikelihood[rolls]
    return torch.FloatTensor(log_likelihoods,requires_grad=False)
def _compute_likelihood_numerator(self, loglikelihoods, states):
    """Computes numerator of likelihood function for a given sequence.

    We'll iterate over the sequence of states and compute the sum 
    of the relevant transition cost with the log likelihood of the observed
    roll. 
    Input:
    loglikelihoods: torch Variable. Matrix of n_obs x n_states. 
    i,j entry is loglikelihood of observing roll i given state j
    states: sequence of labels
    Output:
    score: torch Variable. Score of assignment. 
    """
    prev_state = self.n_states
    score = Variable(torch.Tensor([0]))
    for index, state in enumerate(states):
        score += self.transition[state, prev_state] + loglikelihoods[index, state]
        prev_state = state
    return score

def _compute_likelihood_denominator(self, loglikelihoods):
    """Implements the forward pass of the forward-backward algorithm.

    We loop over all possible states efficiently using the recursive
    relationship: alpha_t(j) = \sum_i alpha_{t-1}(i) * L(x_t | y_t) * C(y_t | y{t-1} = i)
    Input:
    loglikelihoods: torch Variable. Same input as _compute_likelihood_numerator.
    This algorithm efficiently loops over all possible state sequences
    so no other imput is needed.
    Output:
    torch Variable. 
    """

    # Stores the current value of alpha at timestep t
    prev_alpha = self.transition[:, self.n_states] + loglikelihoods[0].view(1, -1)

    for roll in loglikelihoods[1:]:
        alpha_t = []

        # Loop over all possible states
        for next_state in range(self.n_states):
            
            # Compute all possible costs of transitioning to next_state
            feature_function = self.transition[next_state,:self.n_states].view(1, -1) +\
            roll[next_state].view(1, -1).expand(1, self.n_states)

        alpha_t_next_state = prev_alpha + feature_function
        alpha_t.append(self.log_sum_exp(alpha_t_next_state))
        prev_alpha = torch.cat(alpha_t).view(1, -1)
    return self.log_sum_exp(prev_alpha)